import os
from tqdm import tqdm
import logging
import time
from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from TableStructureRec.get_table_image import save_poly_image_from_pdf
from utils.title_and_table_change import title_and_table_change
from utils.table_merge.table_md_merge import remove_body_html_tags
from log_config import setup_logging
import time
# 日志初始化
setup_logging()

# 表格结构识别主流程
def table_predict_optimized(infer_result, pdf_path, local_image_dir):
    """
    并行优化：表格图片收集、裁剪、批量推理，HTML写回
    """
    # 阶段一：收集所有表格检测任务
    # gt_time_start = time.time()
    tasks = []
    for page in infer_result._infer_res:
        page_no = page['page_info']['page_no']
        page_size = (
            page['page_info']['width'],
            page['page_info']['height']
        )
        for layout_det in page['layout_dets']:
            if layout_det.get('category_id') == 5:  # 5表示表格
                tasks.append({
                    'layout_det': layout_det,
                    'page_no': page_no,
                    'page_size': page_size,
                })
    # gt_time_end = time.time()
    # logging.info(f"阶段一，表格检测任务收集完成。耗时：{gt_time_end - gt_time_start:.2f}秒")

    # 阶段二：提取并保存表格图片
    for task in tqdm(tasks, desc="提取并保存表格图片"):
        det = task['layout_det']
        img_path, _ = save_poly_image_from_pdf(
            pdf_path,
            task['page_no'],
            det['poly'],
            task['page_size'],
            local_image_dir
        )
        task['img_path'] = img_path

    # 阶段三：批量推理生成HTML
    t2h_start_time = time.time()
    img_paths = [task['img_path'] for task in tasks]
    from pymupdf_text.pmp_imgtable2html import table2html
    # t2h1_end_time = time.time()
    # logging.info(f"阶段三，第一部分。耗时：{t2h1_end_time - t2h_start_time:.2f}秒")
    table_results = table2html(
        img_paths=img_paths,
        pdf_path=pdf_path,
        tasks=tasks,
    )
    # t2h2_end_time = time.time()
    # logging.info(f"阶段三，第二部分。耗时：{t2h2_end_time - t2h1_end_time:.2f}秒")
    for table_result, task in zip(table_results, tasks):
        det = task['layout_det']
        det['html'] = table_result.pred_html
    t2h_end_time = time.time()
    # logging.info(f"阶段三，第三部分。耗时：{t2h_end_time - t2h2_end_time:.2f}秒")
    logging.info(f"阶段三，批量推理生成HTML完成。耗时：{t2h_end_time - t2h_start_time:.2f}秒")
    return infer_result

# def table_predict(infer_result,pdf_path,local_image_dir,ds):
#     # 获取table数据：
#     table_tasks = []
#     for page in infer_result._infer_res:
#         page_no = page['page_info']['page_no']
#         page_size = {
#             'width':page['page_info']['width'],
#             'height':page['page_info']['height']
#         }
#         for layout_det in page['layout_dets']:
#             if layout_det.get('category_id') == 5:  # 5表示表格
#                 table_tasks.append({
#                     'layout_det': layout_det,
#                     'page_no': page_no,
#                     'page_size': page_size,
#                 })
#     # 获取text坐标
#     # 获取单元格坐标
#     ...

def pdf_to_markdown(
    pdf_path: str,
    output_dir: str ="output",
    pre_md_save: bool = False,
    ocr: bool = False,
    table_recognize: bool = False,
    processes_draw: bool = False
):
    """
    从PDF文件提取表格内容，识别结构，输出Markdown文件（含表格HTML），全流程封装。
    Args:
        pdf_path (str): PDF文件路径
        output_dir (str): 输出文件夹路径
        pre_md_save (bool): 是否保存预处理的Markdown文件
        ocr (bool): 是否使用OCR进行文本识别
    Returns:
        final_md_path (str): 最终生成的Markdown文件路径
        markdown路径, 标题变换后内容
    """
    start_time = time.time()
    # 路径准备
    name_without_extension = os.path.basename(pdf_path).split('.')[0]
    logging.info(f'运行{pdf_path:}')
    local_md_dir = output_dir + f"/{name_without_extension}"
    local_image_dir = output_dir + f"/{name_without_extension}/images"
    image_dir = str(os.path.basename(local_image_dir))
    os.makedirs(local_image_dir, exist_ok=True)

    image_writer = FileBasedDataWriter(local_image_dir)
    md_writer = FileBasedDataWriter(local_md_dir)

    # 读取PDF内容
    reader1 = FileBasedDataReader("")
    pdf_bytes = reader1.read(pdf_path)

    # 执行文档初步结构分析
    ds = PymuDocDataset(pdf_bytes)
    
    if ocr:
        # 如果需要OCR，执行OCR识别
        infer_result = ds.apply(doc_analyze, ocr=True)
    else:
        # 否则直接执行文档分析
        infer_result = ds.apply(doc_analyze, ocr=False)

    if  table_recognize:
        # 执行优化后的表格识别流程
        tpo_time_start = time.time()
        logging.info("开始进行自定义表格识别！！！")
        infer_result = table_predict_optimized(
            infer_result,
            pdf_path=pdf_path,
            local_image_dir=local_image_dir
        )
        # infer_result = table_predict(
        #     infer_result,
        #     pdf_path=pdf_path,
        #     local_image_dir=local_image_dir,
        #     ds = ds
        # )
        tpo_time_end = time.time()
        logging.info(f"自定义表格识别完成，耗时：{tpo_time_end - tpo_time_start:.2f}秒")

    # 生成Markdown内容，写入预处理md文件
    if ocr:
        # 如果使用OCR，使用OCR模式生成Markdown
        pipe_result = infer_result.pipe_ocr_mode(image_writer)
    else:
        # 否则使用文本模式生成Markdown
        pipe_result = infer_result.pipe_txt_mode(image_writer)

    # 保持处理前数据
    if pre_md_save:
        md_content = pipe_result.get_markdown(image_dir)
        md_content = remove_body_html_tags(md_content)
        md_writer.write_string(f"{name_without_extension}_pre.md",md_content)
        logging.info(f"标题处理和表格合并前数据保存成功：\n{output_dir}/{name_without_extension}/{name_without_extension}_pre.md")

    if processes_draw:
        ### get model inference result
        model_inference_result = infer_result.get_infer_res()

        ### draw layout result on each page
        pipe_result.draw_layout(os.path.join(local_md_dir, f"{name_without_extension}_layout.pdf"))

        ### draw spans result on each page
        pipe_result.draw_span(os.path.join(local_md_dir, f"{name_without_extension}_spans.pdf"))

        ### get markdown content
        md_content = pipe_result.get_markdown(image_dir)

        ### dump markdown
        pipe_result.dump_md(md_writer, f"{name_without_extension}.md", image_dir)

        ### get content list content
        content_list_content = pipe_result.get_content_list(image_dir)

        ### dump content list
        pipe_result.dump_content_list(md_writer, f"{name_without_extension}_content_list.json", image_dir)

        ### get middle json
        middle_json_content = pipe_result.get_middle_json()

        ### dump middle json
        pipe_result.dump_middle_json(md_writer, f'{name_without_extension}_middle.json')
    # 尝试进行标题与表格样式调整
    try:
        md_content = title_and_table_change(pipe_result, local_image_dir, image_dir)
        logging.info("title_and_table_change 执行成功。")
    except Exception as e:
        logging.error(f"title_and_table_change 执行出错: {e}", exc_info=True)
        md_content = '程序错误，转换失败！！！'

    # 写入最终Markdown文件
    md_writer.write_string(f"{name_without_extension}.md", md_content)
    end_time = time.time()
    logging.info(f"Markdown文件保存成功\n{output_dir}/{name_without_extension}/{name_without_extension}.md")
    logging.info(f"all time:{end_time-start_time}")
    final_md_path = os.path.join(local_md_dir, f"{name_without_extension}.md")
    final_md_path = os.path.abspath(final_md_path)
    return final_md_path, md_content


if __name__ == "__main__":
    pdf_path = "../input_doc/sample/提取自√3.14-国家电力投资集团有限公司2025年度第五期中期票据(能源保供特别债)募集说明书.pdf"
    pdf_md_path, md_text = pdf_to_markdown(
        pdf_path=pdf_path,
        output_dir="../output",
        processes_draw=True,
        table_recognize=False
    )
    print(f"Markdown文件路径: {pdf_md_path}")
    